{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "ae69c074",
   "metadata": {},
   "source": [
    "# 긴 문맥 재정렬(LongContextReorder)\n",
    "\n",
    "모델의 아키텍처와 상관없이, 10개 이상의 검색된 문서를 포함할 경우 성능이 상당히 저하됩니다.\n",
    "\n",
    "간단히 말해, 모델이 긴 컨텍스트 중간에 있는 관련 정보에 접근해야 할 때, 제공된 문서를 무시하는 경향이 있습니다.\n",
    "\n",
    "자세한 내용은 다음 논문을 참조하세요\n",
    "\n",
    "- https://arxiv.org/abs/2307.03172\n",
    "\n",
    "이 문제를 피하기 위해, 검색 후 문서의 순서를 재배열하여 성능 저하를 방지할 수 있습니다.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9283f84e",
   "metadata": {},
   "source": [
    "- `Chroma` 벡터 저장소를 사용하여 텍스트 데이터를 저장하고 검색할 수 있는 `retriever`를 생성합니다.\n",
    "- `retriever`의 `invoke` 메서드를 사용하여 주어진 쿼리에 대해 관련성이 높은 문서를 검색합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1a914221",
   "metadata": {},
   "outputs": [],
   "source": [
    "# API 키를 환경변수로 관리하기 위한 설정 파일\n",
    "from dotenv import load_dotenv\n",
    "\n",
    "# API 키 정보 로드\n",
    "load_dotenv()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0d337cfb",
   "metadata": {},
   "outputs": [],
   "source": [
    "# LangSmith 추적을 설정합니다. https://smith.langchain.com\n",
    "# !pip install langchain-teddynote\n",
    "from langchain_teddynote import logging\n",
    "\n",
    "# 프로젝트 이름을 입력합니다.\n",
    "logging.langsmith(\"CH10-Retriever\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "089f999e",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain_core.prompts import PromptTemplate\n",
    "from langchain_community.document_transformers import LongContextReorder\n",
    "from langchain_community.vectorstores import Chroma\n",
    "from langchain_openai import OpenAIEmbeddings\n",
    "\n",
    "\n",
    "# 임베딩을 가져옵니다.\n",
    "embeddings = OpenAIEmbeddings(model=\"text-embedding-3-small\")\n",
    "\n",
    "texts = [\n",
    "    \"이건 그냥 내가 아무렇게나 적어본 글입니다.\",\n",
    "    \"사용자와 대화하는 것처럼 설계된 AI인 ChatGPT는 다양한 질문에 답할 수 있습니다.\",\n",
    "    \"아이폰, 아이패드, 맥북 등은 애플이 출시한 대표적인 제품들입니다.\",\n",
    "    \"챗GPT는 OpenAI에 의해 개발되었으며, 지속적으로 개선되고 있습니다.\",\n",
    "    \"챗지피티는 사용자의 질문을 이해하고 적절한 답변을 생성하기 위해 대량의 데이터를 학습했습니다.\",\n",
    "    \"애플 워치와 에어팟 같은 웨어러블 기기도 애플의 인기 제품군에 속합니다.\",\n",
    "    \"ChatGPT는 복잡한 문제를 해결하거나 창의적인 아이디어를 제안하는 데에도 사용될 수 있습니다.\",\n",
    "    \"비트코인은 디지털 금이라고도 불리며, 가치 저장 수단으로서 인기를 얻고 있습니다.\",\n",
    "    \"ChatGPT의 기능은 지속적인 학습과 업데이트를 통해 더욱 발전하고 있습니다.\",\n",
    "    \"FIFA 월드컵은 네 번째 해마다 열리며, 국제 축구에서 가장 큰 행사입니다.\",\n",
    "]\n",
    "\n",
    "\n",
    "# 검색기를 생성합니다. (K는 10으로 설정합니다)\n",
    "retriever = Chroma.from_texts(texts, embedding=embeddings).as_retriever(\n",
    "    search_kwargs={\"k\": 10}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8887d361",
   "metadata": {},
   "source": [
    "검색기에 쿼리를 입력하여 검색을 수행합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2871aeaf",
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"ChatGPT에 대해 무엇을 말해줄 수 있나요?\"\n",
    "\n",
    "# 관련성 점수에 따라 정렬된 관련 문서를 가져옵니다.\n",
    "docs = retriever.invoke(query)\n",
    "docs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b50f6b95",
   "metadata": {},
   "source": [
    "`LongContextReorder` 클래스의 인스턴스인 `reordering`을 생성합니다.\n",
    "\n",
    "- `reordering.transform_documents(docs)`를 호출하여 문서 목록 `docs`를 재정렬합니다.\n",
    "  - 덜 관련된 문서는 목록의 중간에 위치하고, 더 관련된 문서는 시작과 끝에 위치하도록 재정렬됩니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae93a9a1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 문서를 재정렬합니다\n",
    "# 덜 관련된 문서는 목록의 중간에 위치하고 더 관련된 요소는 시작/끝에 위치합니다.\n",
    "reordering = LongContextReorder()\n",
    "reordered_docs = reordering.transform_documents(docs)\n",
    "\n",
    "# 4개의 관련 문서가 시작과 끝에 위치하는지 확인합니다.\n",
    "reordered_docs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c6292e32",
   "metadata": {},
   "source": [
    "## Context Reordering을 사용하여 질의-응답 체인 생성\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f2ea0a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_docs(docs):\n",
    "    return \"\\n\".join([doc.page_content for i, doc in enumerate(docs)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4a645502",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(format_docs(docs))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c9b07e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "def format_docs(docs):\n",
    "    return \"\\n\".join(\n",
    "        [\n",
    "            f\"[{i}] {doc.page_content} [source: teddylee777@gmail.com]\"\n",
    "            for i, doc in enumerate(docs)\n",
    "        ]\n",
    "    )\n",
    "\n",
    "\n",
    "def reorder_documents(docs):\n",
    "    # 재정렬\n",
    "    reordering = LongContextReorder()\n",
    "    reordered_docs = reordering.transform_documents(docs)\n",
    "    combined = format_docs(reordered_docs)\n",
    "    print(combined)\n",
    "    return combined"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9af013ad",
   "metadata": {},
   "source": [
    "재정렬된 문서를 출력합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a193bc14",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 재정렬된 문서를 출력\n",
    "_ = reorder_documents(docs)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "703afe03",
   "metadata": {},
   "outputs": [],
   "source": [
    "from langchain.prompts import ChatPromptTemplate\n",
    "from operator import itemgetter\n",
    "from langchain_openai import ChatOpenAI\n",
    "from langchain_core.output_parsers import StrOutputParser\n",
    "from langchain_core.runnables import RunnableLambda\n",
    "\n",
    "# 프롬프트 템플릿\n",
    "template = \"\"\"Given this text extracts:\n",
    "{context}\n",
    "\n",
    "-----\n",
    "Please answer the following question:\n",
    "{question}\n",
    "\n",
    "Answer in the following languages: {language}\n",
    "\"\"\"\n",
    "\n",
    "# 프롬프트 정의\n",
    "prompt = ChatPromptTemplate.from_template(template)\n",
    "\n",
    "# Chain 정의\n",
    "chain = (\n",
    "    {\n",
    "        \"context\": itemgetter(\"question\")\n",
    "        | retriever\n",
    "        | RunnableLambda(reorder_documents),  # 질문을 기반으로 문맥을 검색합니다.\n",
    "        \"question\": itemgetter(\"question\"),  # 질문을 추출합니다.\n",
    "        \"language\": itemgetter(\"language\"),  # 답변 언어를 추출합니다.\n",
    "    }\n",
    "    | prompt  # 프롬프트 템플릿에 값을 전달합니다.\n",
    "    | ChatOpenAI(model=\"gpt-4o-mini\")  # 언어 모델에 프롬프트를 전달합니다.\n",
    "    | StrOutputParser()  # 모델의 출력을 문자열로 파싱합니다.\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89f08a63",
   "metadata": {},
   "source": [
    "`question` 에 쿼리를 입력하고 `language` 에 언어를 입력합니다.\n",
    "\n",
    "- 재정렬된 문서의 검색 결과도 확인합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9249067c",
   "metadata": {},
   "outputs": [],
   "source": [
    "answer = chain.invoke(\n",
    "    {\"question\": \"ChatGPT에 대해 무엇을 말해줄 수 있나요?\", \"language\": \"KOREAN\"}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1f24418a",
   "metadata": {},
   "source": [
    "답변을 출력합니다.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ef4b288",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(answer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py-test",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
